import os
from sklearn.model_selection import train_test_split
import shutil

# 源数据集根目录
source_data_root = r'data\source_data'

# 目标训练集和验证集根目录
target_root = r'data\flower_data'
os.makedirs(target_root, exist_ok=True)

# 获取所有图像的文件路径和标签
all_images = []
all_labels = []

for class_folder in os.listdir(source_data_root):
    class_path = os.path.join(source_data_root, class_folder)
    if os.path.isdir(class_path):
        for image_file in os.listdir(class_path):
            image_path = os.path.join(class_path, image_file)
            all_images.append(image_path)
            all_labels.append(class_folder)

# 划分数据集
train_images, val_images, train_labels, val_labels = train_test_split(
    all_images, all_labels, test_size=0.2, random_state=42, stratify=all_labels
)

# 创建目标训练集和验证集目录结构
for split in ['train', 'val']:
    target_split_dir = os.path.join(target_root, split)
    os.makedirs(target_split_dir, exist_ok=True)

    for class_folder in os.listdir(source_data_root):
        class_path = os.path.join(source_data_root, class_folder)
        if os.path.isdir(class_path):
            target_class_dir = os.path.join(target_split_dir, class_folder)
            os.makedirs(target_class_dir, exist_ok=True)


# 复制图像到目标目录
def copy_images(image_paths, source_root, target_root):
    for image_path in image_paths:
        relative_path = os.path.relpath(image_path, source_root)
        target_path = os.path.join(target_root, relative_path)
        shutil.copyfile(image_path, target_path)


copy_images(train_images, source_data_root, os.path.join(target_root, 'train'))
copy_images(val_images, source_data_root, os.path.join(target_root, 'val'))
